import cv2 
import os
import numpy as np
import sys
sys.path.append('././')
from kinematic import *
from UART import *
# 标定板 9*6 25mm
chess_board_x_num = 10 - 1
chess_board_y_num = 7 - 1
chess_board_len = 25

# 相机内参矩阵和畸变系数
cam_mtx = np.array( [[670.92263044  ,      0.        ,   358.82135372],
                    [  0.           ,    673.8728403 ,   224.19994634],
                    [  0.           ,      0.        ,     1.        ]], dtype=np.float64)

cam_dist = np.array([-0.19770613, -0.19343703 ,-0.00058339, -0.00063087 , 0.07889203], dtype=np.float64)
def get_RT_from_chessboard(img, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist):
    '''
    #用来从棋盘格图片得到相机外参
    :param img_path: 读取图片路径
    :param chess_board_x_num: 棋盘格x方向格子数
    :param chess_board_y_num: 棋盘格y方向格子数
    :param chess_board_len: 单位棋盘格长度,mm
    :param cam_mtx: 相机内参
    :param cam_dist: 相机畸变参数
    :return: 相机外参
    '''

    
    # termination criteria
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
    # 标定板世界坐标
    objp = np.zeros((chess_board_y_num*chess_board_x_num,3), np.float32)
    for m in range(chess_board_y_num):
        for n in range(chess_board_x_num):
            objp[m*chess_board_x_num + n] = [n*chess_board_len, m*chess_board_len, 0]
    # print('obj:\n',objp)
    # Arrays to store object points and image points from all the images.
    objpoints = [] # 3d point in real world space
    imgpoints = [] # 2d points in image plane.


    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # Find the chess board corners
    ret, corners = cv2.findChessboardCorners(gray, (chess_board_x_num, chess_board_y_num), None)

    # If found, add object points, image points (after refining them)
    if ret == True:
        objpoints.append(objp)
        corners2 = cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), criteria)
        imgpoints.append(corners2)
        draw_p=cv2.drawChessboardCorners(img, (chess_board_x_num, chess_board_y_num), corners, ret)
        # cv2.imwrite('calibration/handeye_clib/drawChessboardCorners/'+str(count)+'.jpg', draw_p)
        # count+=1
        # Draw and display the corners
        # cv2.namedWindow('img', 0)
        # cv2.resizeWindow('img', int(2560*0.6), int(1440*0.6))
        # cv2.drawChessboardCorners(img, (chess_board_x_num, chess_board_y_num), corners2, ret)
        # cv2.imshow('img', img)
        # cv2.waitKey(10)
        # input("请输入任意字符继续：")
        retval, rvec, tvec = cv2.solvePnP(objpoints[0], imgpoints[0], cam_mtx, cam_dist)
        # print(rvec.reshape((1,3)))
        # 反投影误差
        total_error = 0
        imgpoints2, _ = cv2.projectPoints(objpoints[0], rvec, tvec, cam_mtx, cam_dist)
        error = cv2.norm(imgpoints[0], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
        total_error += error
        print("SolvePnP error: ", total_error / len(objpoints))

        RT = np.column_stack(((cv2.Rodrigues(rvec))[0],tvec))
        RT = np.row_stack((RT, np.array([0, 0, 0, 1])))
        


        return True,RT,draw_p
    return False,None,None

def test():
    # 打开摄像头
    uart=UART('COM3',115200)
    uart.open_port()

    import time
    time.sleep(1)
    uart.send_command([0, 96,-36, -96, 0],time='1000')
    time.sleep(1)
    uart.set_grapper_pwm(1900)
    capture = cv2.VideoCapture(0)

    if not capture.isOpened():
        print("无法打开摄像头.")
        return
    #标定结果
    RT_cam_to_end = np.load('calibration/handeye_clib/eye_to_end.npy')

    RT_end_to_base=forward_kinematic(np.deg2rad([0, 96,-36, -96, 0]))
    
    while True:
        # 读取视频帧
        ret, frame = capture.read()

        if not ret:
            print("无法获取视频帧.")
            break
        

        ret,RT_chess_to_cam,draw_p = get_RT_from_chessboard(frame, chess_board_x_num, chess_board_y_num, chess_board_len, cam_mtx, cam_dist)
        if ret:
            RT_chess_to_base=RT_end_to_base@RT_cam_to_end@RT_chess_to_cam
          
            cv2.putText(draw_p, "chessTobase:"+"{:.2f}".format(RT_chess_to_base[0,3])+','+"{:.2f}".format(RT_chess_to_base[1,3])+','+"{:.2f}".format(RT_chess_to_base[2,3]), (10, 100), cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 0), 1)
                  # 在窗口中显示帧    
            cv2.imshow('draw_p', draw_p) 
        else:
            print('找不到棋盘')
        a=cv2.waitKey(1)
        # 检查按键，按下 'q' 键退出循环
        if a ==27:
            break

    # 释放摄像头资源
    capture.release()
    cv2.destroyAllWindows()


test()